package utils

import "encoding/json"

type Map[K comparable, V any] struct {
	GoroutineLock
	M map[K]V
}

func NewMap[K comparable, V any]() Map[K, V] {
	return Map[K, V]{
		M: make(map[K]V),
	}
}

func (m *Map[K, V]) String() string {
	encodeStr, _ := json.Marshal(m.M)
	return string(encodeStr)
}

func (m *Map[K, V]) MarshalJSON() ([]byte, error) {
	return json.Marshal(m.M)
}

func (m *Map[K, V]) Clear() {
	m.Lock()
	defer m.Unlock()
	m.M = make(map[K]V)
}

func (m *Map[K, V]) IsNil() bool {
	return m.M == nil
}

func (m *Map[K, V]) Put(key K, value V) {
	m.Lock()
	defer m.Unlock()
	m.M[key] = value
}

func (m *Map[K, V]) PutWithGid(key K, value V, goroutineID int64) {
	m.LockWithGid(goroutineID)
	defer m.UnlockWithGid(goroutineID)
	m.M[key] = value
}

// Get 获取键对应的值
func (m *Map[K, V]) Get(key K) (V, bool) {
	m.RLock()
	defer m.RUnlock()
	value, exists := m.M[key]
	return value, exists
}

// GetWithGid 获取键对应的值
func (m *Map[K, V]) GetWithGid(key K, goroutineID int64) (V, bool) {
	m.LockWithGid(goroutineID)
	defer m.UnlockWithGid(goroutineID)
	value, exists := m.M[key]
	return value, exists
}

// Remove 移除指定键的键值对
func (m *Map[K, V]) Remove(key K) {
	m.Lock()
	defer m.Unlock()
	delete(m.M, key)
}

func (m *Map[K, V]) RemoveWithGid(key K, goroutineID int64) {
	m.LockWithGid(goroutineID)
	defer m.UnlockWithGid(goroutineID)
	delete(m.M, key)
}

func (m *Map[K, V]) Len() int {
	m.RLock()
	defer m.RUnlock()
	return len(m.M)
}

func (m *Map[K, V]) RLockRange(f func(key K, value V)) {
	m.RLock()
	defer m.RUnlock()
	for k, v := range m.M {
		f(k, v)
	}
}

func (m *Map[K, V]) LockRange(f func(key K, value V)) {
	m.Lock()
	defer m.Unlock()
	for k, v := range m.M {
		f(k, v)
	}
}

func (m *Map[K, V]) LockRangeWithGid(f func(key K, value V) bool, goroutineID int64) {
	m.LockWithGid(goroutineID)
	defer m.UnlockWithGid(goroutineID)
	for k, v := range m.M {
		needReturn := f(k, v)
		if needReturn {
			return
		}
	}
}

func (m *Map[K, V]) Keys() []K {
	m.RLock()
	defer m.RUnlock()
	keys := make([]K, 0, len(m.M))
	for k := range m.M {
		keys = append(keys, k)
	}
	return keys
}
